#include "eytzinger_aabb_sdf.h"
#include "parallel_for.h"

template <
  typename Derivedp,
  typename Func,
  typename DerivedB,
  typename Derivedleaf
>
IGL_INLINE void igl::eytzinger_aabb_sdf(
  const Eigen::MatrixBase<Derivedp> & p,
  const Func & primitive,
  const Eigen::MatrixBase<DerivedB> & B1,
  const Eigen::MatrixBase<DerivedB> & B2,
  const Eigen::MatrixBase<Derivedleaf> & leaf,
  typename Derivedp::Scalar & f)
{
  using Scalar = typename Derivedp::Scalar;
  using RowVectorS = Eigen::Matrix<Scalar, 1, Derivedp::ColsAtCompileTime>;
  // https://iquilezles.org/articles/distfunctions/
  const auto box_sdf = [&p](const RowVectorS& b1, const RowVectorS& b2) -> Scalar
  {
    // half diagonal
    const RowVectorS b = 0.5*(b2 - b1);
    // relative position
    const RowVectorS r = p - 0.5*(b1 + b2);
    const RowVectorS q = r.cwiseAbs() - b;
    const Scalar t1 = q.cwiseMax(0.0).squaredNorm();
    const Scalar t2 = std::min(q.maxCoeff(),0.0);
    return std::sqrt(t1 + t2);
  };

  // List of current active nodes
  std::vector<std::pair<int,Scalar>> active;
  active.reserve(50);
  active.emplace_back(0,box_sdf(B1.row(0), B2.row(0)));
  f = std::numeric_limits<Scalar>::infinity();
  while(!active.empty())
  {
    const auto [i, box_f] = active.back();
    active.pop_back();
    if(box_f >= f){ continue; }
    if(leaf(i) >= 0)
    {
      f = std::min(f,primitive(leaf(i)));
    }else
    {
      const int left_i = 2*i + 1;
      const int right_i = 2*i + 2;
      const Scalar left_f = box_sdf(B1.row(left_i), B2.row(left_i));
      const Scalar right_f = box_sdf(B1.row(right_i), B2.row(right_i));
      if(left_f < right_f)
      {
        active.emplace_back(right_i, right_f);
        active.emplace_back(left_i, left_f);
      }else
      {
        active.emplace_back(left_i, left_f);
        active.emplace_back(right_i, right_f);
      }
    }
  }
}

template <
  typename DerivedP,
  typename Func,
  typename DerivedB,
  typename Derivedleaf,
  typename DerivedS
>
IGL_INLINE void igl::eytzinger_aabb_sdf(
  const Eigen::MatrixBase<DerivedP> & P,
  const Func & primitive,
  const Eigen::MatrixBase<DerivedB> & B1,
  const Eigen::MatrixBase<DerivedB> & B2,
  const Eigen::MatrixBase<Derivedleaf> & leaf,
  Eigen::PlainObjectBase<DerivedS> & S)
{
  using Scalar = typename DerivedS::Scalar;
  S.resize(P.rows());
  //for(int i = 0; i < P.rows(); i++)
  igl::parallel_for(P.rows(), [&](const int i)
  {
    const Eigen::Matrix<Scalar,1,DerivedP::ColsAtCompileTime> p = P.row(i);
    const std::function<Scalar(const int)> primitive_i = [&](const int j)
    {
      return primitive(p,j);
    };
    igl::eytzinger_aabb_sdf(p,primitive_i,B1,B2,leaf,S(i));
  },1000);
}

#ifdef IGL_STATIC_LIBRARY
// Explicit template instantiation
// generated by autoexplicit.sh
template void igl::eytzinger_aabb_sdf<Eigen::Matrix<double, -1, 3, 1, -1, 3>, std::function<double (Eigen::Matrix<double, 1, 3, 1, 1, 3> const&, int)>, Eigen::Matrix<double, -1, 3, 1, -1, 3>, Eigen::Matrix<int, -1, 1, 0, -1, 1>, Eigen::Matrix<double, -1, 1, 0, -1, 1>>(Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>> const&, std::function<double (Eigen::Matrix<double, 1, 3, 1, 1, 3> const&, int)> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 1, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, 1, 0, -1, 1>> const&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 1, 0, -1, 1>>&);
template void igl::eytzinger_aabb_sdf<Eigen::Matrix<double, 1, 3, 1, 1, 3>, std::function<double (int)>, Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<int, -1, 1, 0, -1, 1>>(Eigen::MatrixBase<Eigen::Matrix<double, 1, 3, 1, 1, 3>> const&, std::function<double (int)> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, 1, 0, -1, 1>> const&, Eigen::Matrix<double, 1, 3, 1, 1, 3>::Scalar&);
template void igl::eytzinger_aabb_sdf<Eigen::Matrix<double, -1, 3, 0, -1, 3>, std::function<double (Eigen::Matrix<double, 1, 3, 1, 1, 3> const&, int)>, Eigen::Matrix<double, -1, 3, 0, -1, 3>, Eigen::Matrix<int, -1, 1, 0, -1, 1>, Eigen::Matrix<double, -1, 1, 0, -1, 1>>(Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, std::function<double (Eigen::Matrix<double, 1, 3, 1, 1, 3> const&, int)> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<double, -1, 3, 0, -1, 3>> const&, Eigen::MatrixBase<Eigen::Matrix<int, -1, 1, 0, -1, 1>> const&, Eigen::PlainObjectBase<Eigen::Matrix<double, -1, 1, 0, -1, 1>>&);
#endif
